package org.github.manx98.jlibhac.tools.fssystem.ncautils;

import org.github.manx98.jlibhac.common.Utilities;
import org.github.manx98.jlibhac.common.keys.KeySet;
import org.github.manx98.jlibhac.common.keys.MissingKeyException;
import org.github.manx98.jlibhac.common.math.IByteUtil;
import org.github.manx98.jlibhac.common.math.Uint64;
import org.github.manx98.jlibhac.crypto.Aes;
import org.github.manx98.jlibhac.fs.*;
import org.github.manx98.jlibhac.fs.fsa.IFileSystem;
import org.github.manx98.jlibhac.fssystem.*;
import org.github.manx98.jlibhac.spl.AccessKey;
import org.github.manx98.jlibhac.tools.fssystem.*;
import org.github.manx98.jlibhac.tools.fssystem.romfs.RomFsFileSystem;

import javax.crypto.ShortBufferException;
import java.io.IOException;

/**
 * @author Manx98
 * @version 1.0
 * @since 2022/10/3 15:17
 */
public class Nca {
    private final KeySet keySet;
    private final IStorage baseStorage;
    public final NcaHeader header;
    private IStorage nca0TransformedBody;
    private byte[] Nca0KeyArea;
    private static final String[] KakNames = new String[]{"application", "ocean", "system"};

    public Nca(KeySet keySet, IStorage storage) throws IOException {
        this.keySet = keySet;
        this.baseStorage = storage;
        this.header = new NcaHeader(keySet, storage);
    }

    public static int getSectionIndexFromType(NcaSectionType type, NcaContentType contentType) {
        int index = tryGetSectionIndexFromType(type, contentType);
        if (index < 0) {
            throw new IllegalArgumentException("NCA does not contain this section type: {type=" + type + ", contentType=" + contentType + "}");
        }
        return index;
    }

    private static int tryGetSectionIndexFromType(NcaSectionType type, NcaContentType contentType) {
        if (type == NcaSectionType.Code && contentType == NcaContentType.Program) {
            return 0;
        }
        if (type == NcaSectionType.Data && contentType == NcaContentType.Program) {
            return 1;
        }

        if (type == NcaSectionType.Logo && contentType == NcaContentType.Program) {
            return 2;
        }
        if (type == NcaSectionType.Data) {
            return 0;
        }
        return -1;
    }

    public boolean sectionExists(NcaSectionType type) {
        int index = tryGetSectionIndexFromType(type, header.getContentType());
        if (index == -1) {
            return false;
        }
        return header.isSectionEnabled(index);
    }

    /**
     * @param type NcaSectionType 常量
     */
    public IFileSystem openFileSystemByType(NcaSectionType type, IntegrityCheckLevel integrityCheckLevel) throws IOException {
        return openFileSystem(getSectionIndexFromType(type), integrityCheckLevel);
    }

    public IFileSystem openFileSystem(int index, IntegrityCheckLevel integrityCheckLevel) throws IOException {
        IStorage storage = openStorage(index, integrityCheckLevel);
        NcaFsHeader header = getFsHeader(index);
        return openFileSystem(storage, header);
    }

    private IFileSystem openFileSystem(IStorage storage, NcaFsHeader header) throws IOException {
        switch (header.getFormatType()) {
            case Pfs0:
                return new PartitionFileSystem(storage);
            case Romfs:
                return new RomFsFileSystem(storage);
            default:
                throw new IllegalArgumentException("Unknown format type: " + header.getFormatType());
        }
    }

    private IStorage openStorage(int index, IntegrityCheckLevel integrityCheckLevel) throws IOException {
        return openStorage(index, integrityCheckLevel, false);
    }

    private IStorage openStorage(int index, IntegrityCheckLevel integrityCheckLevel, boolean leaveCompressed) throws IOException {
        IStorage rawStorage = openRawStorage(index);
        NcaFsHeader header = getFsHeader(index);

        if (header.getEncryptionType() == NcaEncryptionType.AesCtrEx) {
            return rawStorage.slice(0, header.getPatchInfo().getRelocationTreeOffset());
        }

        IStorage returnStorage = createVerificationStorage(integrityCheckLevel, header, rawStorage);

        if (!leaveCompressed && header.existsCompressionLayer()) {
            returnStorage = openCompressedStorage(header, returnStorage);
        }

        return returnStorage;
    }

    private IStorage openCompressedStorage(NcaFsHeader header, IStorage baseStorage) throws IOException {
        NcaCompressionInfo compressionInfo = header.getCompressionInfo();
        BucketTree.Header bucketTreeHeader = new BucketTree.Header();
        IByteUtil.mapToStruct(bucketTreeHeader, compressionInfo.MetaHeader, 0);
        bucketTreeHeader.verify();
        long nodeStorageSize = CompressedStorage.queryNodeStorageSize(bucketTreeHeader.EntryCount);
        long entryStorageSize = CompressedStorage.queryEntryStorageSize(bucketTreeHeader.EntryCount);
        long tableOffset = compressionInfo.MetaOffset;
        long tableSize = compressionInfo.MetaSize;

        if (entryStorageSize + nodeStorageSize > tableSize) {
            throw new IllegalArgumentException("Entry size " + entryStorageSize + nodeStorageSize + " exceeds table size limit " + tableSize);
        }

        var dataStorage = new ValueSubStorage(baseStorage, 0, tableOffset);
        var nodeStorage = new ValueSubStorage(baseStorage, tableOffset, nodeStorageSize);
        var entryStorage = new ValueSubStorage(baseStorage, tableOffset + nodeStorageSize, entryStorageSize);
        var compressedStorage = new CompressedStorage();
        compressedStorage.initialize(dataStorage,nodeStorage,entryStorage, bucketTreeHeader.EntryCount);

        return new CachedStorage(compressedStorage, 0x4000, 32, true);
    }

    private IStorage createVerificationStorage(IntegrityCheckLevel integrityCheckLevel, NcaFsHeader header, IStorage rawStorage) {
        switch (header.getHashType()) {
            case Sha256:
                return initIvfcForPartitionFs(header.getIntegrityInfoSha256(), rawStorage, integrityCheckLevel, true);
            case Ivfc:
                // The FS header of an NCA0 section with IVFC verification must be manually skipped
                if (this.header.isNca0()) {
                    rawStorage = rawStorage.slice(0x200);
                }
                return initIvfcForRomFs(header.getIntegrityInfoIvfc(), rawStorage, integrityCheckLevel, true);
            default:
                throw new IllegalStateException("Invalid raw storage hash type: " + header.getHashType());
        }
    }

    private IStorage initIvfcForRomFs(NcaFsIntegrityInfoIvfc ivfc, IStorage dataStorage, IntegrityCheckLevel integrityCheckLevel, boolean leaveOpen) {
        var initInfo = new IntegrityVerificationInfo[ivfc.getLevelCount()];

        initInfo[0] = new IntegrityVerificationInfo(
                new MemoryStorage(ivfc.getMasterHash()),
                0,
                IntegrityStorageType.Save);

        for (int i = 1; i < ivfc.getLevelCount(); i++) {
            IvfcLevel levelInfo = ivfc.getLevelInfo(i - 1);
            initInfo[i] = new IntegrityVerificationInfo(
                    dataStorage.slice(levelInfo.Offset, levelInfo.Size),
                    1 << levelInfo.BlockSize,
                    IntegrityStorageType.RomFs
            );
        }
        return new HierarchicalIntegrityVerificationStorage(initInfo, integrityCheckLevel, leaveOpen);
    }

    private IStorage initIvfcForPartitionFs(NcaFsIntegrityInfoSha256 info, IStorage pfsStorage, IntegrityCheckLevel integrityCheckLevel, boolean leaveOpen) {
        IStorage hashStorage = pfsStorage.slice(info.getLevelOffset(0), info.getLevelSize(0));
        IStorage dataStorage = pfsStorage.slice(info.getLevelOffset(1), info.getLevelSize(1));

        var initInfo = new IntegrityVerificationInfo[3];

        // Set the master hash
        initInfo[0] = new IntegrityVerificationInfo(
                // todo Get hash directly from header
                new MemoryStorage(info.getMasterHash()),
                0,
                IntegrityStorageType.PartitionFs
        );

        initInfo[1] = new IntegrityVerificationInfo(
                hashStorage,
                (int) info.getLevelSize(0),
                IntegrityStorageType.PartitionFs
        );

        initInfo[2] = new IntegrityVerificationInfo(dataStorage,
                info.getBlockSize(),
                IntegrityStorageType.PartitionFs
        );
        return new HierarchicalIntegrityVerificationStorage(initInfo, integrityCheckLevel, leaveOpen);
    }

    private IStorage openRawStorage(int index) throws IOException {
        return openRawStorage(index, false);
    }

    private IStorage openRawStorage(int index, boolean openEncrypted) throws IOException {
        if (header.isNca0())
            return openNca0RawStorage(index, openEncrypted);

        IStorage storage = openSectionStorage(index);

        if (header.isEncrypted == openEncrypted) {
            return storage;
        }

        return openDecryptedStorage(storage, index, !openEncrypted);
    }

    private IStorage openDecryptedStorage(IStorage storage, int index, boolean decrypting) throws IOException {
        NcaFsHeader header = getFsHeader(index);
        switch (header.getEncryptionType()) {
            case None:
                return storage;
            case XTS:
                return openAesXtsStorage(storage, index, decrypting);
            case AesCtr:
                return openAesCtrStorage(storage, index, this.header.getSectionStartOffset(index), header.getCounter());
            case AesCtrEx:
                return openAesCtrExStorage(storage, index, decrypting);
            default:
                throw new IllegalArgumentException("Unknown encryption type: " + header.getEncryptionType());
        }
    }

    private IStorage openAesCtrExStorage(IStorage baseStorage, int index, boolean decrypting) {
        throw new UnsupportedOperationException("not implemented");
    }

    private IStorage openAesCtrStorage(IStorage baseStorage, int index, long offset, Uint64 upperCounter) {
        byte[] key = getContentKey(NcaKeyType.AesCtr);
        byte[] counter = Aes128CtrStorage.createCounter(upperCounter, header.getSectionStartOffset(index));
        return new Aes128CtrStorage(baseStorage, key, offset, counter, true);
    }

    private IStorage openAesXtsStorage(IStorage baseStorage, int index, boolean decrypting) {
        int sectorSize = 0x200;

        byte[] key0 = getContentKey(NcaKeyType.AesXts0);
        byte[] key1 = getContentKey(NcaKeyType.AesXts1);

        // todo: Handle xts for nca version 3
        return new Aes128XtsStorage(baseStorage, key0, key1, sectorSize, true, decrypting);
    }

    private NcaFsHeader getFsHeader(int index) throws IOException {
        if (header.isNca0())
            return getNca0FsHeader(index);

        return header.getFsHeader(index);
    }

    private NcaFsHeader getNca0FsHeader(int index) throws IOException {
        // NCA0 stores the FS header in the first block of the section instead of the header
        IStorage bodyStorage = openNca0BodyStorage(false);
        long offset = header.getSectionStartOffset(index) - 0x400;

        byte[] fsHeaderData = new byte[0x200];
        bodyStorage.read(offset, fsHeaderData, 0, fsHeaderData.length);

        return new NcaFsHeader(fsHeaderData);
    }

    private IStorage openSectionStorage(int index) {
        if (!sectionExists(index)) throw new IllegalArgumentException("section does not exist, index=" + index);

        long offset = header.getSectionStartOffset(index);
        long size = header.getSectionSize(index);

        long ncaStorageSize = baseStorage.getSize();

        NcaFsHeader fsHeader = header.getFsHeader(index);

        if (fsHeader.existsSparseLayer()) {
            NcaSparseInfo sparseInfo = fsHeader.getSparseInfo();
            BucketTree.Header header = new BucketTree.Header();
            IByteUtil.mapToStruct(header, sparseInfo.MetaHeader, 0);
            header.verify();
            throw new UnsupportedOperationException("not implemented!");
        }
        if (!isSubRange(offset, size, ncaStorageSize)) {
            throw new IllegalArgumentException(String.format(
                    "Section offset (0x%x) and length (0x%x) fall outside the total NCA length (0x%x).",
                    offset,
                    size,
                    ncaStorageSize
            ));
        }
        return baseStorage.slice(offset, size);
    }

    private boolean sectionExists(int index) {
        return header.isSectionEnabled(index);
    }

    private IStorage openNca0RawStorage(int index, boolean openEncrypted) {
        if (!sectionExists(index)) {
            throw new IllegalStateException("Cannot open Nca0 storage for index [ " + index + " ], section not exist!");
        }

        long offset = header.getSectionStartOffset(index) - 0x400;
        long size = header.getSectionSize(index);

        IStorage bodyStorage = openNca0BodyStorage(openEncrypted);

        long baseSize = bodyStorage.getSize();

        if (!isSubRange(offset, size, baseSize)) {
            throw new IllegalArgumentException(
                    String.format(
                            "Section offset (0x%x) and length (0x%x) fall outside the total NCA length (0x%x).",
                            offset + 0x400,
                            size,
                            baseSize + 0x400
                    )
            );
        }

        return new SubStorage(bodyStorage, offset, size);
    }

    private boolean isSubRange(long startIndex, long subLength, long length) {
        boolean isOutOfRange = startIndex < 0 || startIndex > length || subLength < 0 || startIndex > length - subLength;
        return !isOutOfRange;
    }

    private IStorage openNca0BodyStorage(boolean openEncrypted) {
        // NCA0 encrypts the entire NCA body using AES-XTS instead of
        // using different encryption types and IVs for each section.

        if (openEncrypted == header.isEncrypted) {
            return getRawStorage();
        }

        if (nca0TransformedBody != null) {
            return nca0TransformedBody;
        }

        byte[] key0 = getContentKey(NcaKeyType.AesXts0);
        byte[] key1 = getContentKey(NcaKeyType.AesXts1);

        nca0TransformedBody = new Aes128XtsStorage(getRawStorage(), key0, key1, NcaHeader.HeaderSectorSize, true, !openEncrypted);
        return nca0TransformedBody;
    }

    /**
     * @param ncaKeyType NcaKeyType常量
     */
    private byte[] getContentKey(NcaKeyType ncaKeyType) {
        return header.hasRightsId() ? getDecryptedTitleKey() : getDecryptedKey(ncaKeyType.ordinal());
    }

    private byte[] getDecryptedKey(int index) {
        if (index < 0 || index > 3)
            throw new IllegalArgumentException("key index [" + index + "] out of range, must be between 0 and 3");

        // Handle old NCA0s that use different key area encryption
        if (header.formatVersion == NcaVersion.Nca0FixedKey || header.formatVersion == NcaVersion.Nca0RsaOaep) {
            return IByteUtil.slice(getDecryptedKeyAreaNca0(), 0x10 * index, 0x10);
        }

        int keyRevision = Utilities.getMasterKeyRevision(header.getKeyGeneration());
        byte[] keyAreaKey = keySet.KeyAreaKeys[keyRevision][header.getKeyAreaKeyIndex()].getData();

        if (Utilities.isZeros(keyAreaKey)) {
            String keyName = String.format("key_area_key_%s_%02x", KakNames[header.getKeyAreaKeyIndex()], keyRevision);
            throw new MissingKeyException("Unable to decrypt NCA section: " + keyName);
        }

        byte[] encryptedKey = header.getEncryptedKey(index);
        return Aes.decryptEcb128(encryptedKey, keyAreaKey);
    }

    private byte[] getDecryptedKeyAreaNca0() {
        if (Nca0KeyArea != null)
            return Nca0KeyArea;

        if (header.formatVersion == NcaVersion.Nca0FixedKey) {
            Nca0KeyArea = header.GetKeyArea();
        } else {
            throw new UnsupportedOperationException("Unsupported key area type: " + header.formatVersion);
        }

        return Nca0KeyArea;
    }

    private byte[] getDecryptedTitleKey() {
        int keyRevision = Utilities.getMasterKeyRevision(header.getKeyGeneration());
        byte[] titleKek = keySet.TitleKeks[keyRevision].getData();

        var rightsId = new RightsId(header.getRightsId());
        AccessKey accessKey = keySet.ExternalKeySet.get(rightsId);
        if (accessKey == null) {
            throw new MissingKeyException("Missing NCA title key: " + rightsId);
        }

        if (Utilities.isZeros(titleKek)) {
            String keyName = String.format("titlekek_%02x", keyRevision);
            throw new MissingKeyException("Unable to decrypt title key: " + keyName);
        }

        byte[] encryptedKey = accessKey.getData();
        return Aes.decryptEcb128(encryptedKey, titleKek);
    }

    private IStorage getRawStorage() {
        long ncaSize = baseStorage.getSize();
        return baseStorage.slice(0x400, ncaSize - 0x400);
    }

    private int getSectionIndexFromType(NcaSectionType type) {
        return getSectionIndexFromType(type, header.getContentType());
    }
}
